""" This module provides support for counting hashes for interesting files.
We have the directory with files which are search patterns.
We count SHA1 and MD5 from these files.
"""

import pyflag.FlagFramework as FlagFramework
import pyflag.Reports as Reports
import pyflag.DB as DB
import pyflag.pyflaglog as pyflaglog
import hashlib
import os
import pyflag.conf
import pyflag.DB as DB

config=pyflag.conf.ConfObject()

from pyflag.Scanner import *
from pyflag.ColumnTypes import InodeIDType, IntegerType, FilenameType, ColumnType
from stat import *
import csv


class HashMD5Type(ColumnType):
    def __init__(self, **kwargs):
        ColumnType.__init__(self, name="MD5Hash", column='md5hash', **kwargs)

    def create(self):
        return "`%s` char( 32 ) NOT NULL default ''" % self.column

    def display(self, value, row, result):
        return value.upper()
    

class HashSHA1Type(ColumnType):
    def __init__(self, **kwargs):
        ColumnType.__init__(self, name="SHA1Hash", column='sha1hash', **kwargs)

    def create(self):
        return "`%s` char( 40 ) NOT NULL default ''" % self.column

    def display(self, value, row, result):
        return value.upper()
class IdType(ColumnType):
    def __init__(self, **kwargs):
        ColumnType.__init__(self, name="inode_id", column='inode_id', **kwargs)

    def create(self):
        
        return "`%s` int NOT NULL auto_increment" % self.column    

class FnameType(ColumnType):
    def __init__(self, **kwargs):
        ColumnType.__init__(self, name="Filename", column='name', **kwargs)

    def create(self):
        return "`%s` text default '' " % self.column 


    
class HashWantedFilesHashesCaseTable(FlagFramework.CaseTable):
   
    """ Hash Table - Lists MD5 SHA1 hashes for Wanted files """
    primary= "inode_id"
    name = 'wantedhashes'
    columns = [ [ IdType, {} ],
                [ HashMD5Type, {} ],
                [ HashSHA1Type, {} ],
                [ FnameType,{}]
                ]
    index = ['inode_id', 'sha1hash', 'md5hash']
    # TODO add reference key to wanted hashes id ;p

class HashWantedFilesCaseTable(FlagFramework.CaseTable):
    ## do autoincrement integer as primary key
    primary= 'inode_id'
    index=['md5hash']
    name = 'wantedfiles'
    columns = [ [ InodeIDType, dict() ],
                [ HashMD5Type, {} ],
                [ HashSHA1Type, {} ],
                [ FilenameType, dict() ]
               ]


class WantedFiles(Reports.report):
    """ Shows wanted files found on analysed disk"""
    family = "Hash Analysis"
    name = "Wanted Files"
    order = 1
    
    def form(self,query,result):
        result.case_selector()
    
    def display(self,query,result):
        dbh=self.DBO(query['case'])
        dbh.expire_cache()
        try:
            result.table(
                elements = [InodeIDType(),
                            FilenameType(case=query['case']),
                            HashMD5Type(),
                            HashSHA1Type()],
                table='wantedfiles',
                case=query['case'],
                )
        except DB.DBError,e:
            result.para("Error reading the Hash table. Did you remember to run the HashScan scanner?")
            result.para("Error reported was:")
            result.text(e,style="red")


    
class LoadWantedFiles(Reports.report):
    """ Count and Show Wanted Files Hashes """
    name = "Load Wanted Files"
    family = "Hashes"
    description="This report will give a table for describing what files are searched"
   
    parameters = {"datafile":"any",
                   "final":"any"}


    def __filemd5(self,filename):
        file= open(filename, "rb")
        m= hashlib.md5()
        m.update(file.read())
        return m.hexdigest()
        
        
    def __filesha(self,filename):
        file = open(filename, "rb")
        s= hashlib.sha1()
        s.update(file.read())
        return s.hexdigest()

    def __prepareInsert(self,filename):
        #TODO check directories
        uploaddir= config.UPLOADDIR
        if uploaddir[-1]=='/':
            uploaddir=uploaddir[1:-1]
        normalized=  uploaddir + os.path.normpath(filename)
        
        sha = self.__filesha(normalized)
        md5= self.__filemd5(normalized)
        filename= filename.replace("\"","\\\"")
        
        insert=" insert into wantedhashes(md5hash, sha1hash, name) "
        insert+= "values ( '%s' , '%s'  , \"%s\" ) "  % (md5 , sha, filename)         

        logmsg= "normalized: " + normalized + " insert: "+ insert + "\n\n\n"
        pyflaglog.log(7,logmsg)
        return insert
               

    def form(self, query, result):
        try:
            self.DBO(query['case']).expire_cache()
            
            result.case_selector()
            result.fileselector("Wanted file or wanted files directory", 'datafile')
            result.checkbox('Click here to confirm','final','ok')
            

        except (KeyError, RuntimeError):
            pass
        
    def analyse(self, query):
        dbh=self.DBO(query['case'])

        filenames= query.getarray("datafile")

        for f in filenames:
            statinfo= os.stat(f)[ST_MODE]
            if S_ISDIR(statinfo):
                for root, dirs, files in os.walk(f):
                    for name in files:
                        insert= self.__prepareInsert(os.path.join(root,name))
                        dbh.execute(insert)
            else:
                insert=self.__prepareInsert(f)
                dbh.execute(insert)
        dbh.fetch()

        pyflaglog.log(7, "analyse invocation")

    def progress(self, query, result):
        result.heading("Loading Wanted Hashes")
        
    def display(self,query,result):
   
        result.heading("Wanted Hashes ")
        result.para("Successfully stored wanted hashes")
                
        
                
        try:
            result.table(
                elements = [ FnameType(),
                             HashMD5Type(),
                             HashSHA1Type()
                             ],
                table='wantedhashes',
                case=query['case'],
              )
            pyflaglog.log(7,"wanted files display invoked")

        except DB.DBError,e:
            result.para("Error reading the wantedhashes table.")
            result.para("Error reported was:")
            result.text(e,style="red")


        
        
        return result
    

        
            
    def reset(self, query):
        pass
    
    
    
class ImportHashFromCSV(Reports.report):
    parameters = {'filename':'any', 'final':'string'}
    name = 'Load Wanted Hashes (CSV)'
    family = 'Hashes'
    description = 'Import hash code (md5 and sha1) from CSV file format'
    order = 10
    __count = 0
    __message = ''
    
    
    def form(self, query, result):
	result.fileselector('Select CSV files', 'filename')
	result.checkbox('Click here when finished', 'final', 'ok')
	return result
	

    def display(self, query, result):
	
	if self.__message == '':
	    result.text('Import success')
	else:
	    result.text(self.__message)
	
	return result
	
	
    def analyse(self, query):
	filenames = query.getarray('filename')
	query.clear('filename')
	
	try:
	    self.__dbh = DB.DBO(query['case'])
	
	    # Read data from CSV files
	    for f in filenames:
		query['filename'] = f
		self.__readFromCSV(f)	
	except DB.DBError, e:
	    print e
	    self.__message += e + '\n'


    def progress(self, query, result):
	result.heading('Please wait...')
	result.text('Read data from file: %s\n' % query['filename'])
	result.text('Total: ' + str(self.__count))
	return result


    # Read data from CSV files
    # Format CSV file: "sha1","md5","name"	    
    def __readFromCSV(self, filename):
	reader = csv.reader(open(filename, 'rb') , delimiter=",")
	
	try:
	    for row in reader:
                if len(row)>0:
                    md5sum = row[1]; sha1sum = row[0]; #name = row[2]
                    self.__count += 1

                    # No md5 sum
                    if md5sum == '':
                        md5sum = '00000000000000000000000000000000'
                    
                    # No sha1 sum
                    if sha1sum == '':
                        sha1sum = '0000000000000000000000000000000000000000'

                    self.__dbh.execute("insert into wantedhashes(md5hash, sha1hash) values ('%s', '%s')" % (md5sum, sha1sum) )

        except csv.Error, e:
	    print 'File %s, line %d: %s' % (filename, reader.line_num, e)    
	    self.__message += e + '\n'


class WantedFilesScan(GenScanFactory):
    """ Scan filesystem finding wanted files """                

    default = True
    #depends = ['TypeScan']
    group = "GeneralForensics"
    
    def __init__(self,fsfd):
        GenScanFactory.__init__(self, fsfd)
      

    class Scan(ScanIfType):
        def __init__(self, inode,ddfs,outer,factories=None,fd=None):
            BaseScanner.__init__(self, inode,ddfs,outer,factories, fd=fd)
            self.m = hashlib.md5()
            self.s= hashlib.sha1()
            self.type = None
            self.length = 0

        def process(self, data,metadata=None):
            #self.boring(metadata,data)
            self.m.update(data)
            self.s.update(data)
            self.length+=len(data)

        def finish(self):
                        
            dbh=DB.DBO(self.case)
            digestmd5 = self.m.hexdigest()
            digestsha= self.s.hexdigest()
 
            dbh.execute("select inode_id from wantedhashes where md5hash like '%s' or \
                                sha1hash like '%s' limit 1", digestmd5, digestsha)

           # print "%d /n/n/n/n/n\n\n\n",self.length
            wantedcounter=dbh.fetch()
            inode_id = self.fd.lookup_id()
           
            #todo add sth like filename
            if wantedcounter:
                #pyflaglog.log(8, "aa wanted found")
                #dbh.execute("insert into wantedfiles(md5hash, sha1hash) values('%s' , '%s' )" , digestmd5, digestsha   )
                dbh.insert('wantedfiles',
                       inode_id = inode_id,
                       __md5hash = digestmd5,
                       __sha1hash = digestsha
                       )
                
                
         #todo reset scaner  
        
